import re
import json
import argparse

class DecompileParser:
    def __init__(self, decompiled_code_path, inline_disassembly_path, cg_graph_path, output_json_path):
        self.decompiled_code_path = decompiled_code_path
        self.inline_disassembly_path = inline_disassembly_path
        self.cg_graph_path = cg_graph_path
        self.output_json_path = output_json_path
        self.decompile_results = []
        self.decompiled_lines = self.load_decompiled_code()
        self.cg_graph = self.load_cg_graph()

    def load_decompiled_code(self):
        with open(self.decompiled_code_path, 'r', encoding='utf-8') as file:
            return file.readlines()

    def load_cg_graph(self):
        with open(self.cg_graph_path, 'r', encoding='utf-8') as file:
            return json.load(file)

    def get_all_caller_paths(self, current_dst_id, nodes, edges, path=None, visited=None, max_depth=100):
        if path is None:
            path = []
        if visited is None:
            visited = set()
        if len(path) > max_depth:
            return []
        paths = []
        for edge in edges:
            if edge.get('dst_id') == current_dst_id:
                src_id = edge.get('src_id')
                for node in nodes:
                    if node.get('id') == src_id:
                        func_name = node.get('function_name')
                        if src_id in visited:
                            continue
                        new_path = path + [func_name]
                        visited.add(src_id)
                        sub_paths = self.get_all_caller_paths(src_id, nodes, edges, new_path, visited, max_depth)
                        visited.remove(src_id)
                        if sub_paths:
                            paths.extend(sub_paths)
                        else:
                            paths.append(new_path)
        return paths

    def parse_inline_disassembly(self):
        with open(self.inline_disassembly_path, 'r', encoding='utf-8') as file:
            content = file.read()

        entries = re.split(r'\n{2,}', content.strip())
        total_entries = len(entries)
        print(f"总共有 {total_entries} 个条目待处理。")

        for idx, entry in enumerate(entries, start=1):
            lines = entry.strip().split('\n')
            print(f"\n处理条目 {idx}: 行数 = {len(lines)}")
            for i, line in enumerate(lines, start=1):
                print(f"  第 {i} 行: {line}")

            if not lines:
                print(f"条目 {idx} 为空，跳过。")
                continue

            first_line = lines[0]
            inline_status = "no" if "NOT INLINED" in first_line else "yes"

            func_match = re.match(r"// Function '([^']+)' in '([^']+)' - (INLINED|NOT INLINED):?", first_line)
            if not func_match:
                print(f"条目 {idx} - 无法匹配函数信息: {first_line}")
                continue
            name, caller_func_name, status = func_match.groups()

            related_par = []
            call_address = ""
            code_content = ""

            if inline_status == "yes":
                if len(lines) < 5:
                    print(f"条目 {idx} - INLINED 条目行数不足: {entry}")
                    continue
                call_address_match = re.search(r"'0x([0-9A-Fa-f]+):", lines[3])
                if call_address_match:
                    call_address = f"0x{call_address_match.group(1)}"
                else:
                    print(f"条目 {idx} - INLINED 条目中无法找到 call_address: {lines[3]}")
                    continue
                code_content_match = re.search(r"Inline Decompile Code '(.+)'", lines[2])
                if code_content_match:
                    code_content = code_content_match.group(1)
                else:
                    print(f"条目 {idx} - INLINED 条目中无法找到 code_content: {lines[2]}")
                    continue
            else:
                if len(lines) < 3:
                    print(f"条目 {idx} - NOT INLINED 条目行数不足: {entry}")
                    continue
                call_address_match = re.match(r"0x([0-9A-Fa-f]+):", lines[1])
                if call_address_match:
                    call_address = f"0x{call_address_match.group(1)}"
                else:
                    print(f"条目 {idx} - NOT INLINED 条目中无法找到 call_address: {lines[1]}")
                    continue
                code_content_match = re.search(r"Decompiled Line: (.+)", lines[2])
                if code_content_match:
                    code_content = code_content_match.group(1)
                else:
                    print(f"条目 {idx} - NOT INLINED 条目中无法找到 code_content: {lines[2]}")
                    continue

            params_match = re.search(r'\w+\((.+)\)', code_content)
            if params_match:
                params = params_match.group(1)
                related_par = [param.strip().strip('"') for param in params.split(',')]
            else:
                print(f"条目 {idx} - 无法提取相关参数: {code_content}")

            code_info = self.get_code_location_info(code_content)
            if code_info["start_line"] == 0:
                print(f"条目 {idx} - 在 decompiled_code.c 中找不到 code_content: {code_content}")

            caller_func_list = self.get_caller_func_list(caller_func_name)

            # record = {
            #     "inline(IDA)": inline_status,
            #     "name(IDA)": name,
            #     "caller_func_name(IDA)": caller_func_name,
            #     "caller_func_list(angr)": caller_func_list,
            #     "related_par(IDA)": related_par,
            #     "call_address(IDA)": call_address,
            #     "code(IDA)": code_info
            # }
            record = {
                "inline": inline_status,
                "name": name,
                "caller_func_name": caller_func_name,
                "caller_func_list": caller_func_list,
                "related_par": related_par,
                "call_address": call_address,
                "code": code_info
            }
            self.decompile_results.append(record)

        print(f"\n解析完成，共生成 {len(self.decompile_results)} 条记录。")

    def get_code_location_info(self, code_content):
        for idx, line in enumerate(self.decompiled_lines, start=1):
            if code_content in line:
                start_column = line.find(code_content) + 1
                end_column = start_column + len(code_content)
                return {
                    "start_line": idx,
                    "end_line": idx,
                    "start_column": start_column,
                    "end_column": end_column,
                    "code_content": line.strip() + "\n"
                }
        return {
            "start_line": 0,
            "end_line": 0,
            "start_column": 0,
            "end_column": 0,
            "code_content": ""
        }

    def get_caller_func_list(self, caller_func_name):
        nodes = self.cg_graph.get('nodes', [])
        edges = self.cg_graph.get('edges', [])
        
        func_id = None
        for node in nodes:
            if node.get('function_name') == caller_func_name:
                func_id = node.get('id')
                break

        if func_id is None:
            return []

        caller_func_list = []
        current_dst_id = func_id

        # while True:
        #     found = False
        #     for edge in edges:
        #         if edge.get('dst_id') == current_dst_id:
        #             src_id = edge.get('src_id')
        #             for node in nodes:
        #                 if node.get('id') == src_id:
        #                     caller_func_list.append(node.get('function_name'))
        #                     current_dst_id = src_id
        #                     found = True
        #                     break
        #             if found:
        #                 break
        #     if not found:
        #         break

        all_paths = self.get_all_caller_paths(current_dst_id, nodes, edges)
        for path in all_paths:
            caller_func_list.append(path)        
        return caller_func_list

    def save_result(self):
        with open(self.output_json_path, 'w', encoding='utf-8') as json_file:
            json.dump({"sensitive_func": self.decompile_results}, json_file, indent=4, ensure_ascii=False)
        print(f"结果已保存到 {self.output_json_path}")

    def run(self):
        self.parse_inline_disassembly()
        self.save_result()

# if __name__ == "__main__":
#     parser = DecompileParser(
#         decompiled_code_path="/home/lxf/angr/results/latex2rtf/decompiled_code.c",
#         inline_disassembly_path="/home/lxf/angr/results/latex2rtf/latex2rtf_32_inline_disassembly.txt",
#         cg_graph_path="/home/lxf/angr/results/latex2rtf/cg_graph.json",
#         output_json_path="/home/lxf/angr/results/latex2rtf/decompile_result.json"
#     )
#     parser.run()

def main():
    parser = argparse.ArgumentParser(description='Decompile Parser')
    parser.add_argument('--decompiled_code_path', type=str, required=True, help='路径到反编译代码文件')
    parser.add_argument('--inline_disassembly_path', type=str, required=True, help='路径到内联反汇编文件')
    parser.add_argument('--cg_graph_path', type=str, required=True, help='路径到控制流图（CG）JSON文件')
    parser.add_argument('--output_json_path', type=str, required=True, help='输出结果的JSON文件路径')

    args = parser.parse_args()

    decompile_parser = DecompileParser(
        decompiled_code_path=args.decompiled_code_path,
        inline_disassembly_path=args.inline_disassembly_path,
        cg_graph_path=args.cg_graph_path,
        output_json_path=args.output_json_path
    )
    decompile_parser.run()

if __name__ == "__main__":
    main()